package com.hk.core.autoconfigure.data.r2dbc;

import com.hk.commons.util.AssertUtils;
import com.hk.core.data.r2dbc.connection.DynamicRoutingConnectionFactory;
import io.r2dbc.spi.ConnectionFactories;
import io.r2dbc.spi.ConnectionFactory;
import io.r2dbc.spi.ConnectionFactoryOptions;
import io.r2dbc.spi.Option;
import lombok.RequiredArgsConstructor;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.r2dbc.R2dbcProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.data.r2dbc.config.AbstractR2dbcConfiguration;
import org.springframework.util.StringUtils;

import java.util.HashMap;
import java.util.Objects;

/**
 * @author Kevin
 * @date 2023-11-12 17:17
 */
@AutoConfiguration
@RequiredArgsConstructor
@ConditionalOnProperty(prefix = "hk.r2dbc.multi-ds", name = "enabled", havingValue = "true")
@EnableConfigurationProperties(value = {R2dbcMultiDataSourceProperties.class})
public class R2DbcMultiDatabaseAutoConfiguration extends AbstractR2dbcConfiguration {

    private final R2dbcMultiDataSourceProperties properties;

    @Override
    public ConnectionFactory connectionFactory() {
        return dynamicRoutingConnectionFactory();
    }

    // 继承了 AbstractR2dbcConfiguration, 不需要再使用 @Bean 注解了
//    @Bean
//    @Primary
    public DynamicRoutingConnectionFactory dynamicRoutingConnectionFactory() {
        var targetCollectionFactory = new HashMap<>();
        String defaultDbName = null;
        ConnectionFactory defaultConnectionFactory = null;
        var dataSourceMap = properties.getDataSources();
        if (Objects.nonNull(dataSourceMap)) {
            for (var entry : dataSourceMap.entrySet()) {
                var dataSource = entry.getValue();
                var builder = ConnectionFactoryOptions.parse(dataSource.getUrl())
                        .mutate()
                        .option(ConnectionFactoryOptions.USER, dataSource.getUsername())
                        .option(ConnectionFactoryOptions.PASSWORD, dataSource.getPassword());
                String name = determineDatabaseName(dataSource);
                if (StringUtils.hasLength(name)) {
                    builder.option(ConnectionFactoryOptions.DATABASE, name);
                }
                var props = dataSource.getProperties();
                if (Objects.nonNull(props)) {
                    props.forEach((key, value) -> builder.option(Option.valueOf(key), value));
                }
                var connectionFactoryOptions = builder.build();
                var connectionFactory = ConnectionFactories.get(connectionFactoryOptions);
                if (Objects.isNull(defaultConnectionFactory)) {
                    defaultDbName = entry.getKey();
                    defaultConnectionFactory = connectionFactory;
                }
                targetCollectionFactory.put(entry.getKey(), connectionFactory);
            }
        }
        AssertUtils.notNull(defaultConnectionFactory, "Default DataSource must not be null.");
        var connectionFactory = new DynamicRoutingConnectionFactory(defaultDbName, defaultConnectionFactory, targetCollectionFactory);
        connectionFactory.afterPropertiesSet();// 如果没有使用 @Bean 注解，需要手动调用 此方法.
        return connectionFactory;
    }

    private String determineDatabaseName(R2dbcProperties properties) {
        if (properties.isGenerateUniqueName()) {
            return properties.determineUniqueName();
        }
        if (StringUtils.hasLength(properties.getName())) {
            return properties.getName();
        }
        return null;
    }

}
